#include "binary_cross_entropy_grad_custom_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_binary_cross_entropy_grad_custom.h"
#include "tiling/platform/platform_ascendc.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__
void binary_cross_entropy_grad_custom(GM_ADDR p, GM_ADDR y, GM_ADDR grad, GM_ADDR weight, GM_ADDR z,
                                      BinaryCrossEntropyGradCustomTilingData tiling);
#endif
extern void GenerateTilingData(uint8_t *tilingBuf, uint32_t blockDim);

int32_t main(int32_t argc, char *argv[])
{
    constexpr uint32_t BLOCK_DIM = 8;
    uint8_t *tiling = nullptr;
    size_t tilingSize = 17 * sizeof(uint32_t);

#ifdef ASCENDC_CPU_DEBUG
    tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *pHost, *yHost, *gradHost, *weightHost, *zHost;
    uint8_t *pDevice, *yDevice, *gradDevice, *weightDevice, *zDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
#endif
    GenerateTilingData(tiling, BLOCK_DIM);
    uint32_t dataTypeSize = 4;
    size_t ByteSize = reinterpret_cast<BinaryCrossEntropyGradCustomTilingData *>(tiling)->totalLength * dataTypeSize;

#ifdef ASCENDC_CPU_DEBUG
    constexpr uint32_t DATA_BLOCK_BYTE = 32;
    auto allocAlign = [&](size_t n) {
        return (uint8_t *)AscendC::GmAlloc((n + DATA_BLOCK_BYTE - 1) / DATA_BLOCK_BYTE * DATA_BLOCK_BYTE);
    };
    uint8_t *p      = allocAlign(ByteSize);
    uint8_t *y      = allocAlign(ByteSize);
    uint8_t *grad   = allocAlign(ByteSize);
    uint8_t *weight = allocAlign(ByteSize);
    uint8_t *z      = allocAlign(ByteSize);

    ReadFile("./input/input_p.bin",     ByteSize, p,      ByteSize);
    ReadFile("./input/input_y.bin",     ByteSize, y,      ByteSize);
    ReadFile("./input/input_grad.bin",  ByteSize, grad,   ByteSize);
    ReadFile("./input/weight.bin",      ByteSize, weight, ByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(binary_cross_entropy_grad_custom, BLOCK_DIM,
                p, y, grad, weight, z,
                *reinterpret_cast<BinaryCrossEntropyGradCustomTilingData *>(tiling));
/home/ma-user/work/ascendbasic/AscendDLBasics/BinaryCrossEntropyGrad/BinaryCrossEntropyGradKernelLaunch_tiling/input/input_p.bin
    WriteFile("./output/output_z.bin", z, ByteSize);

    AscendC::GmFree((void *)p);
    AscendC::GmFree((void *)y);
    AscendC::GmFree((void *)grad);
    AscendC::GmFree((void *)weight);
    AscendC::GmFree((void *)z);
    AscendC::GmFree((void *)tiling);
#else
    CHECK_ACL(aclrtMallocHost((void **)(&pHost),      ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost),      ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&gradHost),   ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&weightHost), ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost),      ByteSize));

    CHECK_ACL(aclrtMalloc((void **)&pDevice,      ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice,      ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&gradDevice,   ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&weightDevice, ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice,      ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_p.bin",     ByteSize, pHost,      ByteSize);
    ReadFile("./input/input_y.bin",     ByteSize, yHost,      ByteSize);
    ReadFile("./input/input_grad.bin",  ByteSize, gradHost,   ByteSize);
    ReadFile("./input/weight.bin",      ByteSize, weightHost, ByteSize);

    CHECK_ACL(aclrtMemcpy(pDevice,      ByteSize, pHost,      ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(yDevice,      ByteSize, yHost,      ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(gradDevice,   ByteSize, gradHost,   ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(weightDevice, ByteSize, weightHost, ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ACLRT_LAUNCH_KERNEL(binary_cross_entropy_grad_custom)(BLOCK_DIM, stream,
                                                          pDevice, yDevice, gradDevice, weightDevice, zDevice,
                                                          reinterpret_cast<BinaryCrossEntropyGradCustomTilingData *>(tiling));
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, ByteSize, zDevice, ByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, ByteSize);

    CHECK_ACL(aclrtFree(pDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(gradDevice));
    CHECK_ACL(aclrtFree(weightDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(pHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(gradHost));
    CHECK_ACL(aclrtFreeHost(weightHost));
    CHECK_ACL(aclrtFreeHost(zHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}